package httpx

import (
	"bytes"
	"io"
	"net/http"
)

type ResponseWrite struct {
	Writer http.ResponseWriter
	Code   int
	Buf    *bytes.Buffer
	Data   any
}

// NewResponseWrite 实现http.ResponseWriter的自定义响应体，可以在中间件中使用，向下传递数据
func NewResponseWrite(writer http.ResponseWriter, code int, data any) *ResponseWrite {
	var buf bytes.Buffer
	return &ResponseWrite{
		Writer: writer,
		Code:   code,
		Buf:    &buf,
		Data:   data,
	}
}

func (w *ResponseWrite) Write(bs []byte) (int, error) {
	w.Buf.Write(bs)
	return w.Writer.Write(bs)
}
func (w *ResponseWrite) Header() http.Header {
	return w.Writer.Header()
}

func (w *ResponseWrite) WriteHeader(code int) {
	w.Code = code
	w.Writer.WriteHeader(code)
}

// GetBody 获取请求体 可以在中间件中使用
func GetBody(r *http.Request) ([]byte, error) {
	if r.Body != http.NoBody {
		var buf bytes.Buffer
		if _, err := buf.ReadFrom(r.Body); err != nil {
			return nil, err
		}
		_ = r.Body.Close()
		r.Body = io.NopCloser(bytes.NewReader(buf.Bytes()))
		return io.ReadAll(bytes.NewReader(buf.Bytes()))
	}
	return nil, nil
}
